# coding=utf-8

import base64
import os
from collections import Counter
from datetime import datetime

from Crypto.Cipher import AES, PKCS1_OAEP
from Crypto.PublicKey import RSA
from Crypto.Random import get_random_bytes


class Demo1():
    """
    加密和解密目录
    """

    def __init__(self):
        # 公钥
        self.PUBLIC_KEY = "-----BEGIN PUBLIC KEY-----\nMIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAqB5VP0IstcsU9WD+E003\nMJtceWJzOOmyGNexZz/6/86DidVR/PNBns2NvqoEx4xpO//YEDsEyb8n2nT4qgsv\n0kL5iN2CDz+0i3b+0cnYJtbFAB/ZA1kpF7trKKfg/DvWgq7MiX8cUCMnNpH69J5V\nK154n48D594AuIWItpG9oNLSG8fRvyxLlAt9ucGISf4w5WAbVqGs6Xt0GCTYEK80\nKriPChQyJ4JLcgAmLFTbV70BR/JbH5jwAyAXYxmFRK7/MdCGaqlxrchbONDUIcoz\nNPftvneMKiCd6Y2TqX5mivnhPfVBsqR3k758Pfx4tsK6hCXeoqpNOv8viIEsQ/ph\nIQIDAQAB\n-----END PUBLIC KEY-----"
        # 私钥
        self.PRIVATE_KEY = "-----BEGIN ENCRYPTED PRIVATE KEY-----\nMIIFJTBPBgkqhkiG9w0BBQ0wQjAhBgkrBgEEAdpHBAswFAQI1rq6Tsa4c70CAkAA\nAgEIAgEBMB0GCWCGSAFlAwQBAgQQS0mdnSYRjKLD6IdA5WPy/gSCBNC1JDEs6RA5\ndedgV5yKjTRB1D6FdbGm8+iIu1ompskkatIHWXbUQVUA4MdYfrb5jUwKaZ8mVtW/\n5hf9qNyGH7I8XvdUOp0iJzhRI/d2+1G1LzZXV5zbi6EX737uNan+N4Qvdjb7A9cn\noILKFGTbRDSEykLEKGiackrL3O0sjpl8d76Z4KQ6kJBkwzp/XqMfDlp/vOgGhX7D\n7nOTT8gPeYYIYctvoti0YW0gBzk/U3ZUPoXSc51Pt3LOngFeJLfbGy+PHKsgM8cH\ntHhmzBopwaLHfy2pLqxO+FihKdI7sk8ZOWSqNrIhr7MdVCQewDef1SUaRFzLN3Gu\nuNeSasBpxIGjLOD7FLfdzJypw57hqEvd+BA+VBraizi6xk9OY6hN5EMeqhSYw0+M\npp+hVryJZnJAGeZwjh6fjM/MmnB8awm33Cekp82006ozlDYBAzCre6EObPhplZra\nQK36M8ma8yZSeHKp44nu/ZoKhcfDrN/7iKFVLEsltjZqgZiaUqkWiZhHrkNj3qee\n0Y+CVGiU5f4HaSid2IX4GeBT2Rj8OiqQpeKH5m16HeJDEfe6Nskgs3hzhvJ9QiUn\n/07WksosakdOEi1F0rkgwU4dTQMP5djCvRg2Qq4RFyTA00XayBCWWAUwVhsZkvet\nq9oFRAOBvR6qvf+nbURgQ5briHehN/xYPbjIgi7AEHePk2q0zR71ABYW/2kxFtXi\noEply0Rd0/v1+78+Nv8NrajP+CR0uGjA7r6fSLDQQMty7zPKgzw+3cEg+HE6BqU2\nIF0KkE8nrW7WBlrK2iBGpmweFrKt/JXfIEdHNEQTCIOruevHZ83N0TnbkHicMB3x\nX29UdlIYVeJ2J9MzZOd+619hcXB0O5KXXtX2+tBqGKN//bRkNcWvDGpRxd1uWX87\nl+/53r+uG7G0IpLjrZaAOArT5/ho4/sPCv/hL1zYGHqcvXHtUjZYQnGJiEx/i4oY\n0UdfTjxMyJ6kshWskvTVfMxQVpoywrbKEdSyRTgQ+/s52y7HdifpkMBvdmSvOOF3\nHC6N9fV4QjihUtJwc0blr5qEeHBmWYrPoiVAiEHRxbB30ANUNzHaJQ0XdF3sOHyA\nIcZ0ITil9DNKpKAqjJz8jEmQwnQ+PSFAQDBmM4Yp4Wztj3vWDQAkOtTtaqMne6lm\nHva3nYGzGw6ckkgVW2EJffqVf8qVzxVpiNy257gKqQVBnb6Cyla3ZJsHxJ+KCpmG\nv2O6hUxEZdTeuB/VyOPOsviqJ34PttpfEaAK/XjpQCfyWtRoQkaMdrI9XbLtfAom\nhhhREEfJOrz1Ori08y6plkyWuUsQrISs/wwQy4S/007NjmDBriohatkJDf34LHwR\nY8Y6GKr+CaevgGvTJx48M0AQeRHGw7250muKnZ5oUdmfSJT/hLifA/kweKFmg4vd\nkhXnCFTmq2PgJvtZVuEDbWRxVBxuuLrOU5EjUrfPi71lI2aj1jcHvF+Y3uywbP0n\ncMifDDz80w53vwFLFuGJIhytVWjocnhd0KgEyktBaQLUQuTOtzdc/VLxrcTeS+X3\nIPBR9T8f3gz0WcNf5PV4NQEC/PiPKsvNrrSHqIONV7PDPKj7ABjpnVUB2mM8u64E\ncDOOXwDC1ik5Lic8Pdg9SN9hib8LafOQfw==\n-----END ENCRYPTED PRIVATE KEY-----"
        # 密码长度
        self.PASSWORD_LEN = 6

    def writeRSAKeys(self, satus):
        """
        生成秘钥对并缓存
        :param satus: 是否打印秘钥对状态
        :return: none
        """
        code = 'nooneknows'
        key = RSA.generate(2048)
        self.PRIVATE_KEY = key.exportKey(passphrase=code, pkcs=8, protection="scryptAndAES128-CBC")
        self.PUBLIC_KEY = key.publickey().exportKey()
        if satus is True:
            print("私钥：" + str(self.PRIVATE_KEY))
            print("公钥：" + str(self.PUBLIC_KEY))

    def createRSAKeys(self):
        code = 'nooneknows'
        key = RSA.generate(2048)
        encrypted_key = key.exportKey(passphrase=code, pkcs=8, protection="scryptAndAES128-CBC")
        # 私钥
        with open('zmy_rsa', 'wb') as f:
            f.write(encrypted_key)
        # 公钥
        with open('zmy_rsa.pub', 'wb') as f:
            f.write(key.publickey().exportKey())

    def Encrypt(self, filename):
        # 目录列表
        data = ''
        with open(filename, 'rb') as f:
            data = f.read()
        with open(filename, 'wb') as out_file:
            # 收件人秘钥 - 公钥
            # recipient_key = RSA.import_key(open('zmy_rsa.pub').read())
            recipient_key = RSA.import_key(bytes(self.PUBLIC_KEY, encoding="utf8"))
            session_key = get_random_bytes(16)
            # Encrypt the session key with the public RSA key
            cipher_rsa = PKCS1_OAEP.new(recipient_key)
            out_file.write(cipher_rsa.encrypt(session_key))
            # Encrypt the data with the AES session key
            cipher_aes = AES.new(session_key, AES.MODE_EAX)
            ciphertext, tag = cipher_aes.encrypt_and_digest(data)
            out_file.write(cipher_aes.nonce)
            out_file.write(tag)
            out_file.write(ciphertext)

    def Descrypt(self, filename):
        code = 'nooneknows'
        with open(filename, 'rb') as fobj:
            # 导入私钥
            # private_key = RSA.import_key(open('zmy_rsa').read(), passphrase=code)
            private_key = RSA.import_key(bytes(self.PRIVATE_KEY, encoding="utf8"), passphrase=code)
            # 会话密钥， 随机数，消息认证码，机密的数据
            enc_session_key, nonce, tag, ciphertext = [fobj.read(x) for x in (private_key.size_in_bytes(), 16, 16, -1)]
            cipher_rsa = PKCS1_OAEP.new(private_key)
            session_key = cipher_rsa.decrypt(enc_session_key)
            cipher_aes = AES.new(session_key, AES.MODE_EAX, nonce)
            # 解密
            data = cipher_aes.decrypt_and_verify(ciphertext, tag)

        with open(filename, 'wb') as wobj:
            wobj.write(data)

    def RenameFile(self, dir, filename):
        filename_bytes = filename.encode('utf-8')
        filename_bytes_base64 = base64.encodebytes(filename_bytes)

        filename_bytes_base64 = filename_bytes_base64[::-1][1:]
        new_filename = filename_bytes_base64.decode('utf-8') + '.liuxp'

        # print (new_filename)
        print(os.path.join(dir, filename))
        print(os.path.join(dir, new_filename))
        os.rename(os.path.join(dir, filename), os.path.join(dir, new_filename))

    def ReserveFilename(self, dir, filename):
        f = filename
        filename = filename[::-1][6:][::-1]
        filename_base64 = filename[::-1] + '\n'
        filename_bytes_base64 = filename_base64.encode('ascii')  # encode as ASCII
        ori_filename = base64.decodebytes(filename_bytes_base64).decode('utf-8')
        new_filename = ori_filename

        # print(new_filename)
        print(os.path.join(dir, f))
        print(os.path.join(dir, new_filename))
        os.rename(os.path.join(dir, f), os.path.join(dir, new_filename))

    # 解密代码
    def decryptPath(self, rootDir):
        """
        解密代码
        :param rootDir: 解密目录
        :return: none
        """
        count = 0
        list_dirs = os.walk(rootDir)
        for root, dirs, files in list_dirs:
            # 遍历文件，解密并且恢复名字
            for f in files:
                if f.endswith('.liuxp'):
                    filename = os.path.join(root, f)
                    self.Descrypt(filename)
                    self.ReserveFilename(root, f)
                    count += 1
        if count > 0:
            print("解密完成！您已成功解密" + str(count) + "个文件。")
        else:
            print("没有需要解密的文件！")

    # 加密代码
    def encryptPath(self, rootDir):
        """
        加密代码
        :param rootDir: 加密路径
        :return: none
        """
        count = 0
        list_dirs = os.walk(rootDir)
        for root, dirs, files in list_dirs:
            # 遍历文件，加密并且改名
            for f in files:
                filename = os.path.join(root, f)
                self.Encrypt(filename)
                self.RenameFile(root, f)
                count += 1
        if count > 0:
            print("加密完成！您已成功加密" + str(count) + "个文件。")
        else:
            print("文件目录为空，没有需要加密的文件！")


    def checkPasswd(self, passwd):
        passwd = str(passwd)

        # Check password length
        if len(passwd) != self.PASSWORD_LEN:
            return False

        # Get current time
        current_time = datetime.now()

        # Extract month, week, and hour information
        mon_num = str((current_time.month // 10 + current_time.month % 10) % 10)
        wek = str((current_time.weekday() + 1) % 10)
        hou_num = str((current_time.hour // 10 + current_time.hour % 10) % 10)

        # print(mon_num, wek, hou_num)

        # Create a dictionary with required counts for each number
        required_counts = {
            num: sum([num == mon_num, num == wek, num == hou_num])
            for num in set([mon_num, wek, hou_num])
        }

        # Count the occurrences of each unique number in passwd
        digit_counts = Counter(passwd)

        # Check if password contains the required number of each unique value
        return all(digit_counts[num] >= count
                   for num, count in required_counts.items())


if __name__ == '__main__':
    d = Demo1()
    ver = "V1.0.9"
    print("\n【BJT372】 " + ver + "\n")
    print(r"     $$$$$$$\     $$$$$\ $$$$$$$$\  $$$$$$\  $$$$$$$$\  $$$$$$\ ")
    print(r"     $$  __$$\    \__$$ |\__$$  __|$$ ___$$\ \____$$  |$$  __$$ |")
    print(r"     $$ |  $$ |      $$ |   $$ |   \_/   $$ |    $$  / \__/  $$ |")
    print(r"     $$$$$$$\ |      $$ |   $$ |     $$$$$ /    $$  /   $$$$$$  |")
    print(r"     $$  __$$\ $$\   $$ |   $$ |     \___$$\   $$  /   $$  ____/")
    print(r"     $$ |  $$ |$$ |  $$ |   $$ |   $$\   $$ | $$  /    $$ |")
    print(r"     $$$$$$$  |\$$$$$$  |   $$ |   \$$$$$$  |$$  /     $$$$$$$$\\")
    print(r"     \_______/  \______/    \__|    \______/ \__/      \________|")
    print("\n声明：\n此工具作用特殊，为防止不法分子恶意使用，本软件需要激活使用。\n")
    while True:
        password = str(input("<< 请输入激活码："))

        # 设置help帮助彩蛋
        if password == "help":
            print("\n1. 长度为6位字符。 \n"
                  "2. 包含月份两位数字之和（如果和大于两位数取最后一位数）。\n"
                  "3. 包含星期数字\n"
                  "4，包含小时（24h制）两位数字之和\n"
                  "例如：2023-11-12 15:03:56 星期日\n"
                  "           2     6           7 \n"
                  "激活码要求6为字符，包含【2】【6】【7】即可，顺序可以错乱，其他3个字符混淆输入。\n"
                  "有效激活码为：w26h7e  776322 6;_72g\n")
            continue

        # 设置名字彩蛋
        if password == '刘晓鹏' or password == 'LXP' or password == 'lxp':
            print("哦？这都被你发现了，给你点个大大的赞！！！\n欢迎探索其他彩蛋噢！\n")
            continue

        # 设置隐藏永久密码
        elif password == "PGV587" or password == "pgv587":
            break

        # 设置日期时间密码
        t = d.checkPasswd(password)
        # print("调试: checkPasswd 结果:", t)
        if t:
            break
        else:
            print("激活码错误，无法激活!!!   退出请按“0”")

        # 设置输入激活码阶段退出程序
        if password == "0":
            exit()

    print("\n<<<<<<<<<<<<<<<<<<<<<<<<欢迎使用【BJT372】加解密软件>>>>>>>>>>>>>>>>>>>>>>")
    while True:
        print("\n***************************** 菜单编号 ********************************")
        print("     0: 退出 （关闭程序）")
        print("     1: 加密 （加密指定目录下所有文件）")
        print("     2: 解密 （解密指定目录下所有文件）")
        print("     3: 关于 （开发人员信息）")
        print("**********************************************************************\n")
        num = input("<< 请输入菜单编号：")
        if len(num) >= 1:
            # 判断输入是否为数字
            if num.isdigit():
                if num == '1':
                    path = input("<< 请输入需要加密的目录：")
                    d.encryptPath(path)
                elif num == '2':
                    path = input("<< 请输入需要解密的目录：")
                    d.decryptPath(path)
                elif num == '3':
                    print(r"半截塔372加解密软件" + ver)
                    print(r" ______________________________________________________________________")
                    print(r"|                                                                      |")
                    print(r"|            _  __   _  _     ___     ___    _  _    _____             |")
                    print(r"|           | |/ /  | \| |   |_ _|   / __|  | || |  |_   _|            |")
                    print(r"|           | ' <   | .` |    | |   | (_ |  | __ |    | |              |")
                    print(r"|           |_|\_\  |_|\_|   |___|   \___|  |_||_|   _|_|_             |")
                    print(r'|          _|"""""|_|"""""|_|"""""|_|"""""|_|"""""|_|"""""|            |')
                    print(r'|           **0-0-***-0-0-***-0-0-***-0-0-***-0-0-***-0-0**            |')
                    print(r"|______________________________________________________________________|")
                    print(r"                     版权所有 翻版必究 KT出品 必属精品                       ")
                    print(r"BY: 软件于2023年11月12日 11时12分25秒在北京市昌平区半截塔村编写")
                    print(r"                                                     -- 不愿透露姓名的刘先生")
                # elif num == '100':
                #     print("正在重新生成新的秘钥对：")
                #     d.writeRSAKeys(True)
                elif num == '0':
                    break
                else:
                    print("此功能正在建设中...")
            else:
                print("请输入有效的数字编号!!!")
        # print("\n|----------------------------------菜单编号---------------------------------|")
        # print("|                     [1]加密   [2]解密   [3]关于   [0]退出                   |")
        # print("|--------------------------------------------------------------------------|")
